from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional
from contextlib import asynccontextmanager
import redis
import json
import logging
from enum import IntEnum
import os

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("market_api")

# Redis连接
redis_client = None

# 从环境变量获取Redis配置
REDIS_HOST = os.getenv("REDIS_HOST", "127.0.0.1")
REDIS_PORT = int(os.getenv("REDIS_PORT", "9998"))
REDIS_DB = int(os.getenv("REDIS_DB", "0"))
REDIS_PASSWORD = os.getenv("REDIS_PASSWORD", "leex123454")

@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理"""
    # 启动前
    global redis_client
    try:
        redis_client = redis.Redis(
            host=REDIS_HOST,
            port=REDIS_PORT,
            db=REDIS_DB,
            password=REDIS_PASSWORD,
            decode_responses=True,
            socket_timeout=5,
            retry_on_timeout=True
        )
        redis_client.ping()
        logger.info(f"Connected to Redis at {REDIS_HOST}:{REDIS_PORT}")
        yield
    except Exception as e:
        logger.error(f"Failed to connect to Redis: {e}")
        if redis_client:
            redis_client.close()
        raise
    finally:
        # 关闭前
        if redis_client:
            redis_client.close()
            logger.info("Disconnected from Redis")

# 创建FastAPI应用
app = FastAPI(
    title="Market Data API",
    description="提供期货市场数据查询服务",
    version="1.0.0",
    lifespan=lifespan
)

class Period(IntEnum):
    """K线周期枚举"""
    TICK = -1
    M1 = 1
    M3 = 3
    M5 = 5
    M10 = 10
    M15 = 15
    M30 = 30
    H1 = 60
    H2 = 120
    DAY = 0
    WEEK = -7
    MONTH = -30

class TickData(BaseModel):
    """Tick数据模型"""
    code: str
    exchg: str
    price: float
    volume: int
    turnover: float
    open_interest: int
    bid_price: float
    bid_volume: int
    ask_price: float
    ask_volume: int
    time: int

class BarData(BaseModel):
    """K线数据模型"""
    code: str
    exchg: str
    period: int
    open: float
    high: float
    low: float
    close: float
    volume: int
    turnover: float
    open_interest: int
    time: int

def check_redis():
    """检查Redis连接状态"""
    if not redis_client:
        raise HTTPException(
            status_code=503,
            detail="Redis connection not available"
        )
    try:
        redis_client.ping()
    except Exception as e:
        logger.error(f"Redis connection error: {e}")
        raise HTTPException(
            status_code=503,
            detail="Redis connection error"
        )

@app.get("/tick/{symbol}")
async def get_tick_data(symbol: str, start_time: Optional[int] = None, end_time: Optional[int] = None):
    """获取Tick数据
    
    Args:
        symbol: 合约代码(格式: SHFE:au2306)
        start_time: 开始时间戳(可选)
        end_time: 结束时间戳(可选)
    
    Returns:
        List[TickData]: Tick数据列表
    """
    check_redis()
    try:
        exchg, code = symbol.split(":")
        key = f"tick:{exchg}:{code}"
        
        # 构建时间范围条件
        if start_time is None and end_time is None:
            data = redis_client.zrange(key, 0, -1, withscores=True)
        else:
            min_score = "-inf" if start_time is None else start_time
            max_score = "+inf" if end_time is None else end_time
            data = redis_client.zrangebyscore(key, min_score, max_score, withscores=True)
        
        # 解析数据
        result = []
        for item, score in data:
            tick_data = json.loads(item)
            result.append(TickData(**tick_data))
        
        logger.info(f"Retrieved {len(result)} tick records for {symbol}")
        return result
    
    except ValueError:
        raise HTTPException(
            status_code=400,
            detail="Invalid symbol format. Expected format: EXCHANGE:CODE"
        )
    except Exception as e:
        logger.error(f"Error fetching tick data for {symbol}: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/bar/{symbol}/{period}")
async def get_bar_data(
    symbol: str, 
    period: Period,
    start_time: Optional[int] = None,
    end_time: Optional[int] = None
):
    """获取K线数据
    
    Args:
        symbol: 合约代码(格式: SHFE:au2306)
        period: K线周期
        start_time: 开始时间戳(可选)
        end_time: 结束时间戳(可选)
    
    Returns:
        List[BarData]: K线数据列表
    """
    check_redis()
    try:
        exchg, code = symbol.split(":")
        key = f"bar:{exchg}:{code}:{period}"
        
        # 构建时间范围条件
        if start_time is None and end_time is None:
            data = redis_client.zrange(key, 0, -1, withscores=True)
        else:
            min_score = "-inf" if start_time is None else start_time
            max_score = "+inf" if end_time is None else end_time
            data = redis_client.zrangebyscore(key, min_score, max_score, withscores=True)
        
        # 解析数据
        result = []
        for item, score in data:
            bar_data = json.loads(item)
            result.append(BarData(**bar_data))
        
        logger.info(f"Retrieved {len(result)} bar records for {symbol} period {period}")
        return result
    
    except ValueError:
        raise HTTPException(
            status_code=400,
            detail="Invalid symbol format. Expected format: EXCHANGE:CODE"
        )
    except Exception as e:
        logger.error(f"Error fetching bar data for {symbol}: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/symbols")
async def get_symbols():
    """获取所有可用的合约列表"""
    check_redis()
    try:
        # 从Redis获取所有tick key
        keys = redis_client.keys("tick:*")
        symbols = []
        for key in keys:
            _, exchg, code = key.split(":")
            symbols.append(f"{exchg}:{code}")
        return symbols
    except Exception as e:
        logger.error(f"Error fetching symbols: {e}")
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        "market_api:app",  # 使用模块导入字符串
        host="0.0.0.0",
        port=9997,
        log_level="info",
        reload=True,  # 开发模式下启用热重载
        reload_dirs=["./"]  # 监视当前目录的变化
    ) 